/*
 * @Description: 前端里程计算法
 * @Author: Ren Qian
 * @Date: 2020-02-04 18:53:06
 */
#include "lidar_localization/matching/matching.hpp"

#include <pcl/common/transforms.h>
#include <pcl/io/pcd_io.h>
#include "glog/logging.h"

#include "lidar_localization/global_defination/global_defination.h"
#include "lidar_localization/models/registration/ndt_registration.hpp"
#include "lidar_localization/models/cloud_filter/voxel_filter.hpp"
#include "lidar_localization/models/cloud_filter/no_filter.hpp"
#include <iostream>
#include <sstream>
#include <fstream>

namespace lidar_localization {

template<class Matrix>
bool read_matrix(const char* filename, Matrix& matrix){
    std::ifstream in(filename, std::ios::in | std::ios::binary);
    if (!in.good())
    {
        return false;
    }
    typename Matrix::Index rows=0, cols=0;
    in.read((char*) (&rows),sizeof(typename Matrix::Index));
    in.read((char*) (&cols),sizeof(typename Matrix::Index));
    matrix.resize(rows, cols);
    in.read( (char *) matrix.data() , rows*cols*sizeof(typename Matrix::Scalar) );
    in.close();
    return true;
}
template<class datatype>
bool read_vector(const char* filename, std::vector<datatype>& vec){
    std::ifstream in(filename, std::ios::in | std::ios::binary);
    if (!in.good())
    {
        return false;
    }
    int num;
    in.read((char*) (&num),sizeof(int));
    vec.resize(num);
    in.read( (char *) vec.data() , num*sizeof(datatype) );
    in.close();
    return true;
}

Matching::Matching()
    :local_map_ptr_(new CloudData::CLOUD()),
     global_map_ptr_(new CloudData::CLOUD()),
     current_scan_ptr_(new CloudData::CLOUD()) {
    
    InitWithConfig();

    InitGlobalMap();

    ResetLocalMap(0.0, 0.0, 0.0);
}

bool Matching::InitWithConfig() {
    std::string config_file_path = WORK_SPACE_PATH + "/config/matching/matching.yaml";
    YAML::Node config_node = YAML::LoadFile(config_file_path);

    std::cout << "-----------------地图定位初始化-------------------" << std::endl;
    InitDataPath(config_node);
    InitRegistration(registration_ptr_, config_node);
    InitFilter("global_map", global_map_filter_ptr_, config_node);
    InitFilter("local_map", local_map_filter_ptr_, config_node);
    InitFilter("frame", frame_filter_ptr_, config_node);
    InitBoxFilter(config_node);
    InitScanContext(config_node["init"]);

    return true;
}

bool Matching::InitDataPath(const YAML::Node& config_node) {
    map_path_ = config_node["map_path"].as<std::string>();
    return true;
}

bool Matching::InitRegistration(std::shared_ptr<RegistrationInterface>& registration_ptr, const YAML::Node& config_node) {
    std::string registration_method = config_node["registration_method"].as<std::string>();
    std::cout << "地图匹配选择的点云匹配方式为：" << registration_method << std::endl;

    if (registration_method == "NDT") {
        registration_ptr = std::make_shared<NDTRegistration>(config_node[registration_method]);
    } else {
        LOG(ERROR) << "没找到与 " << registration_method << " 相对应的点云匹配方式!";
        return false;
    }

    return true;
}

bool Matching::InitFilter(std::string filter_user, std::shared_ptr<CloudFilterInterface>& filter_ptr, const YAML::Node& config_node) {
    std::string filter_mothod = config_node[filter_user + "_filter"].as<std::string>();
    std::cout << "地图匹配" << filter_user << "选择的滤波方法为：" << filter_mothod << std::endl;

    if (filter_mothod == "voxel_filter") {
        filter_ptr = std::make_shared<VoxelFilter>(config_node[filter_mothod][filter_user]);
    } else if (filter_mothod == "no_filter") {
        filter_ptr = std::make_shared<NoFilter>();
    } else {
        LOG(ERROR) << "没有为 " << filter_user << " 找到与 " << filter_mothod << " 相对应的滤波方法!";
        return false;
    }

    return true;
}

bool Matching::InitBoxFilter(const YAML::Node& config_node) {
    box_filter_ptr_ = std::make_shared<BoxFilter>(config_node);
    return true;
}

bool Matching::InitGlobalMap() {
    pcl::io::loadPCDFile(map_path_, *global_map_ptr_);
    LOG(INFO) << "load global map size:" << global_map_ptr_->points.size();

    local_map_filter_ptr_->Filter(global_map_ptr_, global_map_ptr_);
    LOG(INFO) << "filtered global map size:" << global_map_ptr_->points.size();

    has_new_global_map_ = true;

    return true;
}

bool Matching::ResetLocalMap(float x, float y, float z) {
    std::vector<float> origin = {x, y, z};
    box_filter_ptr_->SetOrigin(origin);
    box_filter_ptr_->Filter(global_map_ptr_, local_map_ptr_);

    registration_ptr_->SetInputTarget(local_map_ptr_);

    has_new_local_map_ = true;

    std::vector<float> edge = box_filter_ptr_->GetEdge();
    LOG(INFO) << "new local map:" << edge.at(0) << ","
                                  << edge.at(1) << ","
                                  << edge.at(2) << ","
                                  << edge.at(3) << ","
                                  << edge.at(4) << ","
                                  << edge.at(5) << std::endl << std::endl;

    return true;
}

bool Matching::Initialize(const Eigen::Matrix4f& init_pose, const CloudData& cloud_data)
{
    if (init_method_ == "gnss")
    {
        return SetGNSSPose(init_pose);
    }
    else if (init_method_ == "scan_context")
    {
        std::vector<int> indices;
        CloudData::CLOUD_PTR scan_cloud_ptr(new CloudData::CLOUD());
        pcl::removeNaNFromPointCloud(*cloud_data.cloud_ptr, *scan_cloud_ptr, indices);
        init_scan_filter_ptr_->Filter(scan_cloud_ptr, scan_cloud_ptr);
        std::pair<int, float> rel = init_scmanager_.relocalize(*scan_cloud_ptr);
        if (rel.first == -1)
        {
            return false;
        }
        else
        {
            int key_frame_index = rel.first;
            // 生成地图
            CloudData::CLOUD_PTR map_cloud_ptr(new CloudData::CLOUD());
            Eigen::Matrix4f map_pose = init_key_frame_poses_[key_frame_index];

            for (int i = key_frame_index - init_extend_frame_num_; i < key_frame_index + init_extend_frame_num_; ++i) {
                std::string file_path = init_key_frames_path_ + "/key_frame_" + std::to_string(i) + ".pcd";
                if (i < 0)
                {
                    continue;
                }
                CloudData::CLOUD_PTR cloud_ptr(new CloudData::CLOUD());
                pcl::io::loadPCDFile(file_path, *cloud_ptr);
                
                Eigen::Matrix4f cloud_pose = init_key_frame_poses_[i];
                pcl::transformPointCloud(*cloud_ptr, *cloud_ptr, cloud_pose);

                *map_cloud_ptr += *cloud_ptr;
            }
            init_map_filter_ptr_->Filter(map_cloud_ptr, map_cloud_ptr);
            init_registration_ptr_->SetInputTarget(map_cloud_ptr);
            Eigen::Matrix4f scan_pose = map_pose * (pcl::getTransformation(0, 0, 0, rel.second, 0, 0)).matrix();
            // 匹配
            Eigen::Matrix4f result_pose = Eigen::Matrix4f::Identity();
            CloudData::CLOUD_PTR result_cloud_ptr(new CloudData::CLOUD());
            init_registration_ptr_->ScanMatch(scan_cloud_ptr, scan_pose, result_cloud_ptr, result_pose);
            // 判断是否有效
            if (init_registration_ptr_->GetFitnessScore() > init_fitness_score_limit_)
            {
                std::cout << "检测成功，初始化失败， 关键帧： "<< key_frame_index
                << ", fitness score: " << init_registration_ptr_->GetFitnessScore() << std::endl;
                return false;
            }
            else
            {
                has_inited_ = true;
                SetInitPose(result_pose);
                std::cout << "初始化成功， 关键帧： "<< key_frame_index
                << ", fitness score: " << init_registration_ptr_->GetFitnessScore() << std::endl;
                return true;
            }
        }
    }
    else
    {
        return false;
    }
    
}

bool Matching::Update(const CloudData& cloud_data, Eigen::Matrix4f& cloud_pose) {
    std::vector<int> indices;
    pcl::removeNaNFromPointCloud(*cloud_data.cloud_ptr, *cloud_data.cloud_ptr, indices);

    CloudData::CLOUD_PTR filtered_cloud_ptr(new CloudData::CLOUD());
    frame_filter_ptr_->Filter(cloud_data.cloud_ptr, filtered_cloud_ptr);

    static Eigen::Matrix4f step_pose = Eigen::Matrix4f::Identity();
    static Eigen::Matrix4f last_pose = init_pose_;
    static Eigen::Matrix4f predict_pose = init_pose_;

    if (!has_inited_) {
        predict_pose = current_gnss_pose_;
    }

    // 与地图匹配
    CloudData::CLOUD_PTR result_cloud_ptr(new CloudData::CLOUD());
    registration_ptr_->ScanMatch(filtered_cloud_ptr, predict_pose, result_cloud_ptr, cloud_pose);
    pcl::transformPointCloud(*cloud_data.cloud_ptr, *current_scan_ptr_, cloud_pose);

    // 更新相邻两帧的相对运动
    step_pose = last_pose.inverse() * cloud_pose;
    predict_pose = cloud_pose * step_pose;
    last_pose = cloud_pose;

    // 匹配之后判断是否需要更新局部地图
    std::vector<float> edge = box_filter_ptr_->GetEdge();
    for (int i = 0; i < 3; i++) {
        if (fabs(cloud_pose(i, 3) - edge.at(2 * i)) > 50.0 &&
            fabs(cloud_pose(i, 3) - edge.at(2 * i + 1)) > 50.0)
            continue;
        ResetLocalMap(cloud_pose(0,3), cloud_pose(1,3), cloud_pose(2,3));
        break;
    }

    return true;
}

bool Matching::SetGNSSPose(const Eigen::Matrix4f& gnss_pose) {
    current_gnss_pose_ = gnss_pose;

    static int gnss_cnt = 0;
    if (gnss_cnt == 0) {
        SetInitPose(gnss_pose);
    } else if (gnss_cnt > 3) {
        has_inited_ = true;
    }
    gnss_cnt ++;
    return true;
}

bool Matching::SetInitPose(const Eigen::Matrix4f& init_pose) {
    init_pose_ = init_pose;
    ResetLocalMap(init_pose(0,3), init_pose(1,3), init_pose(2,3));

    return true;
}

void Matching::GetGlobalMap(CloudData::CLOUD_PTR& global_map) {
    global_map_filter_ptr_->Filter(global_map_ptr_, global_map);
    has_new_global_map_ = false;
}

CloudData::CLOUD_PTR& Matching::GetLocalMap() {
    return local_map_ptr_;
}

CloudData::CLOUD_PTR& Matching::GetCurrentScan() {
    return current_scan_ptr_;
}

bool Matching::HasInited() {
    return has_inited_;
}

bool Matching::HasNewGlobalMap() {
    return has_new_global_map_;
}

bool Matching::HasNewLocalMap() {
    return has_new_local_map_;
}

bool Matching::InitScanContext(const YAML::Node& config_node)
{
    init_method_ = config_node["init_method"].as<std::string>();
    if (init_method_ == "gnss")
    {
        return true;
    }
    else if (init_method_ == "scan_context")
    {
        std::string data_path = config_node["data_path"].as<std::string>();
        if (data_path == "./") {
            data_path = WORK_SPACE_PATH;
        }

        init_key_frames_path_ = data_path + "/slam_data/key_frames";
        init_key_frame_poses_.clear();
        std::string filepath = data_path + "/slam_data/trajectory/optimized.txt";
        std::ifstream myfile(filepath);
        std::string temp;
        while (std::getline(myfile, temp))
        {
            std::stringstream input(temp);
            Eigen::Matrix4f p = Eigen::Matrix4f::Identity();
            for (int i = 0; i < 3; ++i) {
                for (int j = 0; j < 4; ++j) {
                    input >> p(i, j);
                }
            }
            init_key_frame_poses_.push_back(p);
        }
        init_extend_frame_num_ = config_node["extend_frame_num"].as<int>();
        init_fitness_score_limit_ = config_node["fitness_score_limit"].as<float>();
        InitFilter("scan", init_scan_filter_ptr_, config_node);
        InitFilter("map", init_map_filter_ptr_, config_node);
        InitRegistration(init_registration_ptr_, config_node);
        init_scmanager_.LIDAR_HEIGHT = config_node["scan_context"]["LIDAR_HEIGHT"].as<double>(); // lidar height : add this for simply directly using lidar scan in the lidar local coord (not robot base coord) / if you use robot-coord-transformed lidar scans, just set this as 0.
        init_scmanager_.PC_NUM_RING = config_node["scan_context"]["PC_NUM_RING"].as<int>(); // 20 in the original paper (IROS 18)
        init_scmanager_.PC_NUM_SECTOR = config_node["scan_context"]["PC_NUM_SECTOR"].as<int>(); // 60 in the original paper (IROS 18)
        init_scmanager_.PC_MAX_RADIUS = config_node["scan_context"]["PC_MAX_RADIUS"].as<double>(); // 80 meter max in the original paper (IROS 18)
        init_scmanager_.PC_UNIT_SECTORANGLE = 360.0 / double(init_scmanager_.PC_NUM_SECTOR);
        init_scmanager_.PC_UNIT_RINGGAP = init_scmanager_.PC_MAX_RADIUS / double(init_scmanager_.PC_NUM_RING);
        init_scmanager_.NUM_EXCLUDE_RECENT = config_node["scan_context"]["NUM_EXCLUDE_RECENT"].as<int>(); // simply just keyframe gap, but node position distance-based exclusion is ok. 
        init_scmanager_.NUM_CANDIDATES_FROM_TREE = config_node["scan_context"]["NUM_CANDIDATES_FROM_TREE"].as<int>(); // 10 is enough. (refer the IROS 18 paper)
        init_scmanager_.SEARCH_RATIO = config_node["scan_context"]["SEARCH_RATIO"].as<double>(); // for fast comparison, no Brute-force, but search 10 % is okay. // not was in the original conf paper, but improved ver.
        init_scmanager_.SC_DIST_THRES = config_node["scan_context"]["SC_DIST_THRES"].as<double>(); // 0.4-0.6 is good choice for using with robust kernel (e.g., Cauchy, DCS) + icp fitness threshold / if not, recommend 0.1-0.15
        init_scmanager_.TREE_MAKING_PERIOD_ = config_node["scan_context"]["TREE_MAKING_PERIOD"].as<int>(); // i.e., remaking tree frequency, to avoid non-mandatory every remaking, to save time cost / in the LeGO-LOAM integration, it is synchronized with the loop detection callback (which is 1Hz) so it means the tree is updated evrey 10 sec. But you can use the smaller value because it is enough fast ~ 5-50ms wrt N.
        for (int i = 0; i < init_key_frame_poses_.size(); ++i)
        {
            Eigen::MatrixXd temp;
            std::string file_path = init_key_frames_path_ + "/scan_context_" + std::to_string(i) + "_";
            read_matrix((file_path + "polarcontexts").c_str(), temp);
            init_scmanager_.polarcontexts_.push_back(temp);
            read_matrix((file_path + "polarcontext_invkeys").c_str(), temp);
            init_scmanager_.polarcontext_invkeys_.push_back(temp);
            read_matrix((file_path + "polarcontext_vkeys").c_str(), temp);
            init_scmanager_.polarcontext_vkeys_.push_back(temp);
            std::vector<float> vt;
            read_vector((file_path + "polarcontext_invkeys_mat").c_str(), vt);
            init_scmanager_.polarcontext_invkeys_mat_.push_back(vt);
        }
        init_scmanager_.treereconstruction();
        return true;
    }
    else
    {
        init_method_ = "gnss";
        return true;
    }
}
}